﻿#pragma once

enum Color {
	RED,
	BLACK
};

template <class T>
struct RBTreeNode {
	RBTreeNode<T>* _left;
	RBTreeNode<T>* _right;
	RBTreeNode<T>* _parent;
	T _data;
	Color _col;

	RBTreeNode(const T& data)
		:_left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _data(data)
		, _col(RED) {
	}
};

//迭代器
template <class T, class Ref, class Ptr>
struct __TreeIterator {
	typedef RBTreeNode<T> Node;
	typedef __TreeIterator<T, Ref, Ptr> Self;
	Node* _node;

	__TreeIterator(Node* node)
		:_node(node) {
	}

	Ref operator*() { return _node->_data; }
	Ptr operator->() { return &_node->_data; }

	Self& operator++() {
		//it.begin()返回的是中序的第一个节点
		//++就是中序遍历的顺序
		//1. it的右不为空，下一个就是其右子树的最左节点
		//2. it的右为空，往上找在其父左的节点(A)，A的父节点就是下一个
		//沿着该节点的路径往上面找孩子是其父的左的祖先节点

		//_node不用担心为空，因为正常使用迭代器它不会为空
		if (_node->_right) { //右不为空。
			Node* cur = _node->_right;
			while (cur->_left) cur = cur->_left;
			_node = cur;
		}
		else { //右为空
			Node* cur = _node;
			Node* parent = cur->_parent;
			while (parent && cur == parent->_right) { //cur是其父的右就继续，不是才停下来
				cur = parent;
				parent = parent->_parent;
			}
			_node = parent;
		}
		return *this;
	}

	Self& operator--() {
		//++是中序 左 根 右，所以--就是中序倒过来右 根 左
		//1.左不为空，左子树的最右节点
		//2.左为空，沿着该节点的路径往上面找孩子是其父的右的那个祖先
		if (_node->_left) {
			Node* cur = _node->_left;
			while (cur->_right) cur = cur->_right;
			_node = cur;
		}
		else {
			Node* cur = _node;
			Node* parent = cur->_parent;
			while (parent && cur == parent->_left) {
				cur = parent;
				parent = parent->_parent;
			}
			_node = parent;
		}
		return *this;
	}

	bool operator!=(const Self& s) { return _node != s._node; }
	bool operator==(const Self& s) { return _node == s._node; }
};


//第二个T实际存储的数据类型，节点中 _data 的具体类型（set存K，map存pair<K, V>）。

//class KeyOfT是为了取T _data 里的key(主要是针对map，因为map的T是pair<K, T>)
//set -> RBTree<K, K, SetKeyOfT> _t;
//map -> RBTree<K, pair<const K, T>, MapKeyOfT> _t;
//红黑树提供一个通用的键值提取接口（KeyOfT），而set和map自行定义如何提取键值
//对于map来说，KeyOfT只是取模版参数T(pair<const K, T>)中K的值(但没有类型)
template <class K, class T, class KeyOfT>
struct RBTree {
	typedef RBTreeNode<T> Node;
public:
	typedef __TreeIterator<T, T&, T*> iterator;
	typedef __TreeIterator<T, const T&, const T*> const_iterator;

	iterator begin() {
		Node* cur = _root;
		while (cur && cur->_left) cur = cur->_left;
		return iterator(cur);
	}
	iterator end() { return iterator(nullptr); }

	const_iterator begin() const {
		Node* cur = _root;
		while (cur && cur->_left) cur = cur->_left;
		return const_iterator(cur);
	}
	const_iterator end() const { return const_iterator(nullptr); }

	//pair<iterator, bool> Insert(const T& data) {
	//set迭代器都是const迭代器，所以传Node*初始化const迭代器
	//迭代器都是用Node*初始化的，所以传Node*可行
	pair<Node*, bool> Insert(const T& data) {
		if (_root == nullptr) { //如果插入节点是根
			_root = new Node(data);
			_root->_col = BLACK;
			return make_pair(_root, true);
		}
		//不是根，找插入位置。记录插入节点的父节点方便链接
		Node* parent = nullptr, * cur = _root;
		KeyOfT kot;
		//data是什么并不确定，KeyOfT是个仿函数类模板，可以取data里面的key
		while (cur) {
			if (kot(data) > kot(cur->_data)) {
				parent = cur;
				cur = cur->_right;
			}
			else if (kot(data) < kot(cur->_data)) {
				parent = cur;
				cur = cur->_left;
			}
			else
				return make_pair(cur, false);
			//insert 返回值规范
			//first: 指向已插入或已存在元素的迭代器
			//second: 是否插入成功（true=新插入，false=已存在）
			//所以插入失败不能返回空，要返回已存在key的位置
		}
		//新增黑色节点会影响所有路径
		//新增红色节点只会影响父节点
		cur = new Node(data);
		Node* newnode = cur;//记录新插入节点是为了最后返回
		cur->_col = RED;
		//插入节点链接在其父节点的左或右
		if (kot(data) > kot(parent->_data)) {
			parent->_right = cur;
			cur->_parent = parent;
		}
		else {
			parent->_left = cur;
			cur->_parent = parent;
		}
		//检查是否需要变色
		//cur的父节点为空(即根节点)，或其父节点不是红色就不用更新
		while (parent && parent->_col == RED) {
			Node* grandfather = parent->_parent;
			if (parent == grandfather->_left) { //通过g找到u。p为g的左孩子
				Node* uncle = grandfather->_right;
				//u存在且为红(cur为红，p为红，g为黑)
				if (uncle && uncle->_col == RED) { //u存在且为红
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					//继续往上更新，去找上一个红节点
					cur = grandfather;//grandfather有可能为根，所以父节点可能不存在，所以需要在循环条件中判断
					parent = cur->_parent;
				}
				else { //u不存在或u为黑(cur为红，p为红，g为黑)
					//p为g的左孩子，cur为p的左孩子，则进行右单旋转。(这里p已经是g的左)
					if (cur == parent->_left) {
						RotateR(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else { //p为g的左孩子，cur为p的右孩子。(这里p已经是g的左)
						RotateL(parent);
						RotateR(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					break;
				}
			}
			else { //parent == grandfather->_right;
				Node* uncle = grandfather->_left;
				if (uncle && uncle->_col == RED) {
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
					parent = cur->_parent;
				}
				else {
					if (cur == parent->_right) {
						RotateL(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else { //cur == parent->_left
						RotateR(parent);
						RotateL(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					break;
				}
			}
		}
		//根节点必须为黑，所以不用在上面的过程中考虑grandfather为根时怎么变色的问题
		_root->_col = BLACK;
		return make_pair(newnode, true);//返回新插入节点
	}

	void RotateL(Node* parent) {
		//parent是平衡因子为2的节点
		Node* subR = parent->_right;
		Node* subRL = subR->_left;
		Node* parentParent = parent->_parent;//parent不一定是整棵树的根
		//旋转完后，subR是子树的根，还需要处理subR的parent

		//旋转
		parent->_right = subRL;
		subR->_left = parent;
		//上方还未处理完，根节点没有确定，各节点的parent还未更新
		parent->_parent = subR;
		if (subRL) subRL->_parent = parent;

		//处理根节点
		if (_root == parent) { //parent是整棵树的根
			_root = subR;
			subR->_parent = nullptr;
		}
		else { //parent不一定是整棵树的根，旋转完后，还需要处理subR的parent
			if (parent == parentParent->_left) parentParent->_left = subR;
			else parentParent->_right = subR;
			subR->_parent = parentParent;
		}
	}

	void RotateR(Node* parent) {
		Node* subL = parent->_left, * subLR = subL->_right, * parentParent = parent->_parent;

		//旋转
		parent->_left = subLR;
		subL->_right = parent;
		//处理各节点的父节点
		parent->_parent = subL;
		if (subLR) subLR->_parent = parent;

		//处理根节点
		if (_root == parent) {
			_root = subL;
			subL->_parent = nullptr;
		}
		else {
			if (parent == parentParent->_left) parentParent->_left = subL;
			else parentParent->_right = subL;
			subL->_parent = parentParent;
		}
	}

	Node* Find(const K& key) {
		Node* cur = _root;
		while (cur) {
			if (key > kot(cur->_data)) cur = cur->_right;
			else if (key < kot(cur->_data)) cur = cur->_left;
			else return cur;
		}
		return nullptr;
	}

	void InOrder() {
		_InOrder(_root);
		cout << endl;
	}

	bool IsBalance() {
		if (_root == nullptr) return true;
		if (_root->_col == RED) return false;
		int blacknum = 0;//根节点到当前节点的黑色节点数量
		//blacknum只能传值，不能传引用。引用求得是整棵树所有黑色节点

		//将最左路径节点作为参考值传入Check和其他路径比较
		int refVal = 0;
		Node* cur = _root;
		while (cur) {
			if (cur->_col == BLACK) ++refVal;
			cur = cur->_left;
		}
		return Check(_root, blacknum, refVal);
	}

	int Height() { return _Height(_root); }
	size_t Size() { return _Size(_root); }
private:
	Node* _root = nullptr;

	void _InOrder(Node* root) {
		if (root == nullptr)
			return;
		_InOrder(root->_left);
		cout << kot(root->_data) << " ";
		_InOrder(root->_right);
	}

	bool Check(Node* root, int blacknum, const int refVal) {
		if (root == nullptr) {
			//cout << blacknum << endl;
			//走到空意味着求出了一条路径的黑色节点数量
			if (refVal != blacknum) {
				cout << "存在黑色节点数量不相等的路径" << endl;
				return false;
			}
			return true;
		}
		if (root->_col == RED && root->_parent->_col == RED) {
			cout << "有连续的红色节点" << endl;
			return false;
		}
		if (root->_col == BLACK) ++blacknum;
		return Check(root->_left, blacknum, refVal)
			&& Check(root->_right, blacknum, refVal);
	}

	int _Height(Node* root) {
		if (root == nullptr) return 0;
		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);

		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}

	size_t _Size(Node* root) {
		if (root == nullptr) return 0;
		return _Size(root->_left) + _Size(root->_right) + 1;
	}
};